# 导入相关模块
import torch
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
import torch.nn as nn
from torchsummary import summary
import torch.optim as optim
from sklearn.model_selection import train_test_split
import numpy as np
import pandas as pd
import time

# 选择设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


# todo: 构建张量数据集对象
def build_dataset():
    # pandas读取csv文件
    data = pd.read_csv(r'./手机价格预测.csv')
    # print(len(data))
    # print(data.head(3).to_string())
    # 获取x和y iloc属性
    x, y = data.iloc[:, :-1], data.iloc[:, -1]
    # print(x)
    # print(y)
    # 数据集分割 训练集和测试集
    x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.8, random_state=88)
    # 将x和y中的数据转换成张量对象数据集
    # x_train.values: 获取数组对象
    train_dataset = TensorDataset(torch.tensor(x_train.values, dtype=torch.float32), torch.from_numpy(y_train.values))
    test_dataset = TensorDataset(torch.tensor(x_test.values, dtype=torch.float32), torch.from_numpy(y_test.values))
    # print('train_dataset--->', train_dataset[0])
    # 返回训练集 测试集张量对象数据集, 特征个数, 类别(标签)个数
    # print('x.shape[1]--->', x.shape[1])
    # print('len(y.unique()--->', len(y.unique()))
    return train_dataset, test_dataset, x.shape[1], len(y.unique())


# todo: 2-构建神经网络模型
class MobileNet(nn.Module):
    # 2-1 定义init方法
    def __init__(self, input_dim, output_dim):
        super().__init__()
        # 定义第一层隐层
        # in_features: 样本特征个数
        self.linear1 = nn.Linear(in_features=input_dim, out_features=128)
        # 定义第二层隐层
        self.linear2 = nn.Linear(in_features=128, out_features=256)
        # 定义输出层
        # out_features: 标签类别个数
        self.out = nn.Linear(in_features=256, out_features=output_dim)

    # 2-2 定义forward方法
    def forward(self, x):
        # print('x--->', x.shape, x)
        # 第一层隐层计算 线性计算+relu激活计算
        x = torch.relu(self.linear1(x))
        # 第二层隐藏计算 线性计算+relu激活计算
        x = torch.relu(self.linear2(x))
        # 输出层计算  线性计算, 后续softmax计算在多分类交叉熵损失函数中实现
        output = self.out(x)
        return output


# todo: 3-训练模型
def train():
    # 获取张量数据集对象
    train_dataset, test_dataset, input_dim, output_dim = build_dataset()
    # 获取数据加载器对象
    train_dataloader = DataLoader(dataset=train_dataset,
                                  batch_size=8,
                                  shuffle=True,
                                  drop_last=True)
    # 创建网络模型对象, 切换训练模式
    model = MobileNet(input_dim, output_dim).to(device=device)
    model.train()
    # 创建优化器对象
    optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)
    # 创建损失函数对象
    criterion = nn.CrossEntropyLoss()
    # 训练模型, 双层训练, 训练轮次
    epochs = 50
    loss_list = []
    # 外层循环遍历轮次
    for epoch in range(epochs):
        # 获取开始时间
        starttime = time.time()
        # 初始化损失值 0.0
        total_loss = 0.0
        # 初始化批次数 0
        total_iter = 0
        # 内层训练遍历数据加载器对象, 批次
        for train_x, train_y in train_dataloader:
            train_x = train_x.to(device=device)
            train_y = train_y.to(device=device)
            # 模型预测
            y_pre = model(train_x)
            # 计算损失值, 统计总损失
            loss = criterion(y_pre, train_y)
            total_loss += loss.item()
            # 梯度清零
            optimizer.zero_grad()
            # 反向传播
            loss.backward()
            # 参数更新
            optimizer.step()
            # 统计批次数
            total_iter += 1
        # 保存损失值
        loss_list.append(total_loss / total_iter)
        # 打印训练信息
        print('epoch: %4s loss: %.2f, time: %.2fs' % (epoch + 1, total_loss / total_iter, time.time() - starttime))


if __name__ == '__main__':
    build_dataset()
    """
    # todo: 构建张量数据集对象
    train_dataset, test_dataset, input_dim, output_dim = build_dataset()
    # todo: 构建数据集加载器对象
    train_dataloader = DataLoader(dataset=train_dataset, batch_size=8, shuffle=True, drop_last=True)
    # todo: 构建神经网络模型对象
    model = MobileNet(input_dim, output_dim).to(device=device)
    # 切换模型训练模式
    model.train()
    # print('model--->', model)
    # 查看模型参数信息
    # summary(model, input_size=(input_dim,), batch_size=8)
    # todo: 循环遍历数据加载器对象
    # train_x, train_y -> 对数据加载器对象的返回值进行拆包
    for train_x, train_y in train_dataloader:
        print('train_x--->', train_x.shape, train_x)
        print('train_y--->', train_y.shape, train_y)
        train_x = train_x.to(device=device)
        train_y = train_y.to(device=device)
        # 调用模型对象进行预测
        y_pred = model(train_x)
        exit()
        """
    train()
